/**
 * Copyright (C) [2019] [NUS Xtra Computing Group & NAOC-TJU Lab]
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published
 * by the Free Software Foundation, either version 3 of the License, or (at
 * your option) any later version. You may obtain a copy of the License at
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * Author: Hao FU (TJU-NUS Joint Ph.D) <haofu@tju.edu.cn>
 */

#ifndef CPU_ONLY
#ifdef USE_PROF

#ifndef __ASYNC_TRACKER_H__
#define __ASYNC_TRACKER_H__

#include <iostream>
#include <fstream>
#include <sstream>
#include <string>
#include <vector>
#include <thread>
#include <mutex>
#include <map>
#include <algorithm>

#include <cstring>
#include <cstdint>

#include <cuda.h>
#include <cuda_runtime.h>
#include <cupti.h>

#include <boost/shared_ptr.hpp>

// Include CUDA-related macro definitions, e.g. CUDA_CHECK, CUBLAS_CHECK, CUDNN_CHECK and CUPTI_CHECK, etc.
#include "caffe/common.hpp"

/**
 * @brief Macro used to represent the buffer size allocated for kernel recording.
 */
#define BUFFER_SIZE (32 * 1024)
#define ALIGN_SIZE (8)            // @brief Alignment size.
#define NAME_SHORT 64
/**
 * @brief Macro used to align profiling buffer.
 */
#define ALIGN_BUFFER(buffer, align_size)                                \
  ((reinterpret_cast<uintptr_t>(buffer) & (align_size - 1)) ?           \
   (buffer + align_size - (reinterpret_cast<uintptr_t>(buffer) & (align_size - 1))) : buffer)

// Standard built-in class used in AsyncResTracker.
namespace caffe {

  /**
   * @brief Structure used to record the start and end timestamp of a kernel.
   */
  typedef struct Timestamp {
    std::string name_; /**< The kernel name */
    uint64_t start_; /**< The start time of a kernel */
    uint64_t end_; /**< The end time of a kernel */
    uint32_t stream_id_; /**< ID of the stream running this kernel */

    // @brief Default constructor.
    Timestamp() :
      name_(""), start_(0), end_(0), stream_id_(0) {}

    // @brief Copy constructor
    Timestamp(const struct Timestamp& other) {
      this->name_ = other.name_;
      this->start_ = other.start_;
      this->end_ = other.end_;
      this->stream_id_ = other.stream_id_;
    }

    // @brief Move constructor.
    Timestamp(struct Timestamp&& other) :
      name_(std::move(other.name_)),
      start_(std::move(other.start_)),
      end_(std::move(other.end_)),
      stream_id_(std::move(other.stream_id_))
    {}

    // @brief Assignment operator.
    struct Timestamp& operator= (const struct Timestamp& other) {
      if (this != &other) {
        this->name_ = other.name_;
        this->start_ = other.start_;
        this->end_ = other.end_;
        this->stream_id_ = other.stream_id_;
      }

      return *this;
    }

    // @brief Move operator.
    struct Timestamp& operator= (struct Timestamp&& other) {
      if (this != &other) {
        this->name_ = std::move(other.name_);
        this->start_ = std::move(other.start_);
        this->end_ = std::move(other.end_);
        this->stream_id_ = std::move(other.stream_id_);
      }

      return *this;
    }

    // @brief Comparison operator.
    bool operator< (const struct Timestamp& r_val) const {
      return (start_ == r_val.start_) ? (end_ < r_val.end_) : (start_ < r_val.start_);
    }
  } Timestamp_t;

  typedef struct Metric {
    CUdevice device_;
    // The set of event groups to collect for a pass.
    CUpti_EventGroupSet *event_groups_;
    // The current number of events collected in eventIdArray and eventValueArray.
    uint32_t event_idx_;
    // the number of entries in eventIdArray and eventValueArray.
    uint32_t num_events_;
    // Array of event ids.
    CUpti_EventID *event_id_array_;
    // array of event values.
    uint64_t *event_val_array_;

    CUpti_MetricID metric_id_;
    CUpti_MetricValueKind val_kind_;
    char metric_name_[NAME_SHORT];
    double double_val_;
    uint64_t uint64_val_;
    int64_t int64_val_;
    double percent_val_;
    uint64_t throughput_val_;
    uint32_t utilization_val_;
  } Metric_t;

  /**
   * @brief Struct designed to store kernel execution configuration and info.
   */
  typedef struct Kernel {
    string name_; /**< Kernel's name in C++ standard */
    int32_t grid_x_, grid_y_, grid_z_; /**< Grid configuration of a kernel (grid_x_, grid_y_, grid_z_) */
    int32_t block_x_, block_y_, block_z_; /**< Thread block configuration of a kernel (block_x_, block_y_, block_z_) */
    uint32_t context_id_; /**< The ID of the context where the kernel is executing */
    uint32_t device_id_; /**< The ID of the device where the kernel is executing */
    /**
     * The cache configuration used for the kernel. The value is one of the CUfunc_cache enumeration values from cuda.h
     */
    // uint8_t executed_;
    /**
     * This indicates if CU_FUNC_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT was updated for the kernel launch
     */
    uint8_t is_sm_carveout_requested_;
    /**
     * This indicates if the kernel was executed via a regular launch or via a single/multi device cooperative launch
     */
    uint8_t launch_type_;
    uint32_t local_mem_per_thread_; /**< The amount of local memory reserved for each thread, in bytes */
    uint32_t local_mem_total_; /**< The total amount of local memory reserved for the kernel, in bytes */
    /**
     * The timestamp when the kernel is queued up in the command buffer, in ns. A value of CUPTI_TIMESTAMP_UNKNOWN
     * indicates that the queued time could not be collected for the kernel. The timestamp is not collected by default.
     * Use API cuptiActivityEnableLatencyTimestamps() to enable collection
     *
     * Command buffer is a buffer written by CUDA driver to send commands like kernel launch, memory copy etc to the GPU.
     * All launches of CUDA kernels are asynchronous with respect to the host, the host requests the launch by writing
     * commands into the command buffer, then return without checking the GPU's progress.
     */
    uint64_t queued_;
    uint16_t reg_per_thread_; /**< The number of registers required for each thread executing the kernel */
    /**
     * The cache configuration requested by the kernel. The value is one of the CUfunc_cache enumeration values from cuda.h
     */
    // uint8_t requested_;
    /**
     * Shared memory carveout value requested for the function in percentage of the total resource.
     * The value will be updated only if field isSharedMemoryCarveoutRequested is set.
     */
    uint8_t sm_carveout_requested_;
    /**
     * The shared memory configuration used for the kernel. The value is one of the CUsharedconfig enumeration values
     * from cuda.h.
     */
    uint8_t sm_config_;
    uint32_t sm_executed_; /**< Sahred memory size set by the driver */
    int32_t dynamic_sm_; /**< The dynamic shared memory reserved for the kernel, in bytes */
    int32_t static_sm_; /**< The static shared memory allocated for the kernel, in bytes */
    uint32_t stream_id_; /**< The ID fo the stream where the kernel is executing. */
    /**
     * The timestamp when the command buffer containing the kernel launch is submitted to the GPU, in ns.
     * A value of CUPTI_TIMESTAMP_UNKNOWN indicates that the submitted time could not be collected for the
     * kernel. This timestamp is not collected by default. Use API cuptiActivityEnableLatencyTimestamps()
     * to enable collection.
     */
    uint64_t submitted_;

    int32_t sm_total_; /**< Total amount of shared memory allocated for the kernel, include both dynamic and static SM */
    unsigned int invocations_;  /**< Total times that this kernel was invoked. */
    uint64_t duration_;  /**< The kernel's execution time. */
    uint64_t average_exec_time_; /**< The average execution time of a kernel. */

    /**
     * @brief Default constructor.
     */
    Kernel() :
      name_(""), grid_x_(0), grid_y_(0), grid_z_(0), block_x_(0), block_y_(0), block_z_(0),
      context_id_(0), device_id_(0)/* , executed_(0)*/, is_sm_carveout_requested_(0), launch_type_(0),
      local_mem_per_thread_(0), local_mem_total_(0), queued_(0), reg_per_thread_(0), /*requested_(0),*/
      sm_carveout_requested_(0), sm_config_(0), sm_executed_(0), dynamic_sm_(0), static_sm_(0),
      stream_id_(0), submitted_(0), sm_total_(0), invocations_(0), duration_(0), average_exec_time_(0)
    {}

    /**
     * @brief Constructor with a CUpti_ActivityKernel4 argument.
     *
     * @param[in] kernel_info The CUpti_ActivityKernel4 object collected by the CUPTI runtime.
     */
    Kernel(const CUpti_ActivityKernel4& kernel_info) {
      name_ = std::string(kernel_info.name);
      grid_x_ = kernel_info.gridX; grid_y_ = kernel_info.gridY; grid_z_ = kernel_info.gridZ;
      block_x_ = kernel_info.blockX; block_y_ = kernel_info.blockY; block_z_ = kernel_info.blockZ;
      context_id_ = kernel_info.contextId;
      device_id_ = kernel_info.deviceId;
      // executed_ = kernel_info.executed;
      is_sm_carveout_requested_ = kernel_info.isSharedMemoryCarveoutRequested;
      launch_type_ = kernel_info.launchType;
      local_mem_per_thread_ = kernel_info.localMemoryPerThread;
      local_mem_total_ = kernel_info.localMemoryTotal;
      queued_ = kernel_info.queued;
      reg_per_thread_ = kernel_info.registersPerThread;
      // requested_ = kernel_info.requested;
      sm_carveout_requested_ = kernel_info.sharedMemoryCarveoutRequested;
      sm_config_ = kernel_info.sharedMemoryConfig;
      sm_executed_ = kernel_info.sharedMemoryExecuted;
      dynamic_sm_ = kernel_info.dynamicSharedMemory;
      static_sm_ = kernel_info.staticSharedMemory;
      stream_id_ = kernel_info.streamId;
      submitted_ = kernel_info.submitted;

      sm_total_ = dynamic_sm_ + static_sm_;
      invocations_ = 0;
      duration_ = 0;
      average_exec_time_ = 0;
    }

    /**
     * @brief Operator== overload.
     *
     * @param[in] other Second kernel object.
     * @return Return true if two object is identical.
     */
    inline bool operator==(const Kernel& other) const {
      return (name_ == other.name_) and (grid_x_ == other.grid_x_) and (grid_y_ == other.grid_y_) and
        (grid_z_ == other.grid_z_) and (block_x_ == other.block_x_) and (block_y_ == other.block_y_) and
        (block_z_ == other.block_z_) and (context_id_ == other.context_id_) and (device_id_ == other.device_id_)
        and /*(executed_ == other.executed_) and*/ (is_sm_carveout_requested_ == other.is_sm_carveout_requested_) and
        (launch_type_ == other.launch_type_) and (local_mem_per_thread_ == other.local_mem_per_thread_) and
        (local_mem_total_ == other.local_mem_total_) and (queued_ == other.queued_) and
        (reg_per_thread_ == other.reg_per_thread_)/* and (requested_ == other.requested_)*/ and
        (sm_carveout_requested_ == other.sm_carveout_requested_) and (sm_config_ == other.sm_config_) and
        (sm_executed_ == other.sm_executed_) and (dynamic_sm_ == other.dynamic_sm_) and (static_sm_ == other.static_sm_)
        and (stream_id_ == other.stream_id_) and (submitted_ == other.submitted_);
    }
  } Kernel_t;

  /**
   * @brief Node structure used to construct a segment tree.
   */
  typedef struct node {
    int left_child_, right_child_; /**< The left and right child of a tree node */
    int start_, end_; /**< The start and end value that a tree node represents */
    bool c_flag_; /**< Flag used to show whether the current node is covered */

    /**
     * @brief Default constructor.
     */
    node() {
      this->left_child_ = this->right_child_ = this->start_ = this->end_ = 0;
      this->c_flag_ = false;
    }

    /**
     * @brief Overloaded operator=.
     */
    struct node& operator= (const struct node& other) {
      this->left_child_ = other.left_child_;
      this->right_child_ = other.right_child_;
      this->start_ = other.start_;
      this->end_ = other.end_;
      this->c_flag_ = other.c_flag_;

      return *this;
    }
  } SegTreeNode_t, *SegTreeNode_ptr;

  class ActivityTracker;

  /**
   * @brief      Kernel activity record parser.
   *
   * @details    Parse the activity record and store the corresponding kernel name
   *             and kernel runtime configuration.
   *
   * @param[in]  kernel_record    A activity record from cuptiActivityGetNextRecord.
   * @param[in]  tracker   The ActivityTracker object.
   */
  void ParseKernelActivity(CUpti_ActivityKernel4 *kernel_record, ActivityTracker& tracker);

  /**
   * @brief      CUPTI overhead record parser.
   *
   * @details    Parse all CUPTT overhead record.
   *
   * @param[in]  overhead_record    A CUPTI overhead record from cuptiActivityGetNextRecord.
   * @param[in]  tracker   The ActivityTracker object.
   */
  void ParseOverheadActivity(CUpti_ActivityOverhead *overhead_record, ActivityTracker& tracker);

  /**
   * @brief      CUPTI metric record parser.
   *
   * @details    Parse all CUPTT metric record.
   *
   * @param[in]  metric_record    A CUPTI metric record from cuptiActivityGetNextRecord.
   * @param[in]  tracker   The ActivityTracker object.
   */
  void ParseMetricActivity(CUpti_ActivityMetric *metric_record, ActivityTracker& tracker);

  /**
   * @brief      CUPTI activity record parser.
   *
   * @details    Parse all CUPTT activity record and store the corresponding kernel name
   *             and kernel runtime configuration.
   *
   * @param[in]  record    A CUPTI activity record from cuptiActivityGetNextRecord.
   * @param[in]  tracker   The ActivityTracker object.
   */
  void ParseCUptiActivity(CUpti_Activity *record, ActivityTracker& tracker);

  /**
   * @brief ActivityTracker class.
   */
  class ActivityTracker
  {
    friend void ParseKernelActivity(CUpti_ActivityKernel4 *kernel_record, ActivityTracker& tracker);
    friend void ParseOverheadActivity(CUpti_ActivityOverhead *overhead_record, ActivityTracker& tracker);
    friend void ParseMetricActivity(CUpti_ActivityMetric *metric_record, ActivityTracker& tracker);
    friend void ParseCUptiActivity(CUpti_Activity *record, ActivityTracker& tracker);

  public:
    /**
     * @brief Object access method.
     *
     * Singleton asynchronous resource tracker object.
     */
    static ActivityTracker& Get();

    /**
     * @brief Initialize CUPTI settings of a class object.
     *
     * @param[in] prof_type  The profiling type.
     *                       Valid values are DEFAULT, SERIAL and CONCURRENT.
     */
    void Init();

    /**
     * @brief Destroy the CUPTI profiling procedure.
     */
    void Destroy();

    /**
     * @brief      Lock function.
     *
     * @details    Function used to lock a profiler.
     */
    void Lock();

    /**
     * @brief      Unlock function.
     *
     * @details    Function used to unlock a profiler.
     */
    void Unlock();

    /**
     * @brief  Function used to start profiling kernels on GPU device_id.
     *
     * @param[in]  device_id    ID of the device needed to be profiled.
     *                          If device_id is lower than 0, then the
     *                          current device is profiled, otherwise
     *                          device device_id is profiled.
     */
    void Start(int device_id = -1);

    /**
     * @brief Function used to stop profiling kernels.
     */
    void Stop();

    /**
     * @brief Get the vector of kernels recorded.
     *
     * @return Reference to the vector of kernels.
     */
    boost::shared_ptr<std::vector<Kernel_t>> GetKernelsRecorded();
    /**
     * @brief Get the overhead of launching a kernel.
     @
     @ return The overhead of launching a kernel on the current device.
     */
    uint64_t GetKernelLaunchOverhead();

    /**
     * @brief    Build a segment tree.
     *
     * @param[in] seg_tree    Pointer to the root of the segment tree.
     * @param[in] node_id     ID of the available tree node.
     * @param[in] start       The start number of the segment tree.
     * @param[in] end         The end number of the segment tree.
     */
    void TreeBuild(SegTreeNode_ptr seg_tree, int node_id, int start, int end);
    /**
     * @brief    Insert a tree node.
     *
     * @param[in] seg_tree    Pointer to the root of the segment tree.
     * @param[in] current_node    ID of the available tree node.
     * @param[in] startp    The start number of the segment tree.
     * @param[in] end    The end number of the segment tree.
     */
    void TreeInsert(SegTreeNode_ptr seg_tree, int current_node, int start, int end);
    /**
     * @brief    Traverse the segment tree.
     *
     * @param[in] seg_tree    Pointer to the root of the segment tree.
     * @param[in] time_vec    Timestamps of kernels recorded.
     * @param[in] node_id     The current tree node being visited.
     *
     * @return     The total busy time of the GPU device.
     */
    uint64_t TreeTraverse(const SegTreeNode_ptr seg_tree, const vector<uint64_t>& time_vec, int node_id);

    /**
     * @brief    Compute the device occupancy ratio.
     *
     * @param[in] layer_name    Name of the current layer.
     * @param[in] parallel_degree    The current parallel_degree configuration.
     */
    void ComputeOccupancyRatio(const string layer_name, const uint32_t parallel_degree);

    /**
     * @brief    Record the kernel timestamp.
     *
     * This method is used to write kernel timestamps recorded for further analysis.
     *
     * @param[in] filename    Name of the log file.
     * @param[in] timestamp_ptr    Kernel timestamp vector.
     */
    void TimestampLog(const string filename) const;

    /**
     * @brief    Temp buffer release method.
     *
     * Method used to release temporary buffer allocated while analyzing kernels recorded.
     */
    void TempBufRelease();

    inline bool SetBuffOverheadRecord(bool buff_overhead_record) {
      this->buff_overhead_record_ = buff_overhead_record;
      
      return buff_overhead_record_;
    }

    inline bool GetBuffOverheadRecord() {
      return buff_overhead_record_;
    }

  protected:

    unsigned int kernel_counter_; /**< Member variable used to store the number of recorded kernels. */

    bool config_flag_; /**< Flag used to identify whether the ActivityTracker has been initialized. */

    int prof_device_id_; /**< Static member variable used to identify the profiling device */

    uint64_t start_timestamp_; /**< Static start timestamp of the profiling process */
    uint64_t kernel_launch_overhead_; /**< Overhead used to launch a kernel */

    boost::shared_ptr<std::vector<Kernel_t>> kernels_vec_sptr_; /**< Member variable used to maintain the recorded kernels. */

    boost::shared_ptr<std::vector<Timestamp_t>> timestamp_vec_sptr_; /**< Member variable used to store kernel timestamps */

    bool buff_overhead_record_; /**< Flag used to control the recording of the CUPTI buffer allocation. */

    bool init_flag_; /**< Flag for initializing the ActivityTracker. */

    /**
     * Member variable used to represent the next variable tree node ID.
     * Only used to construct the segment tree.
     */
    uint32_t next_node_idx_;
    /**
     * Member variable used to represent the total number of nodes in segment tree seg_tree_.
     */
    uint32_t tree_nodes_count_;
    /**
     * Member variable used to point to the root node of the segment tree seg_tree_.
     */
    SegTreeNode_ptr seg_tree_;

  private:
    struct ActivityTrackerProxy {
      ActivityTrackerProxy () {
        ActivityTracker::Get();
      }
    };

    static ActivityTrackerProxy s_activity_tracker_proxy_;

    // @brief Default constructor
    ActivityTracker();

    // @brief Destructor
    ~ActivityTracker() noexcept;

    // Disable the copy, move and assignment operation on ActivityTracker object.
    DISABLE_COPY_MOVE_AND_ASSIGN(ActivityTracker);
  };

}  // caffe

#endif /* __ASYNC_TRACKER_H__ */

#endif // USE_PROF
#endif // CPU_ONLY
